# class to dispatch tasks to many threads, block it when no free thread available

from multiprocessing import Queue
from threading import Thread
from blockingqueue import BlockingQueue
#from multiprocessing import Queue
import time
from threading import Lock

class TaskRunner:
    def __init__(self):
        pass
    
    def __init__(self, threadPool):
        self.setThreadPool(threadPool)
    
    
    def setThreadPool(self, threadPool):
        assert isinstance(threadPool, ThreadPool), "param is not a ThreadPool object"
        self.threadPool = threadPool
        
    def assign(self, task):
        """
        Assign task to a runner.
        """
        self.threadPool.putTask(self, task)
        
    def process(self, task):
        """
        Override this function and define customized proceduce here.
        """
        pass


class ThreadPool:
    def __init__(self, threadNumber, taskBufSize):
        self.kill = object()
        self.threadNum = threadNumber
        #self.taskQueue = Queue(taskBufSize)
        self.taskQueue = BlockingQueue(taskBufSize)
        self.threads = [PoolThread(self) for i in range(threadNumber)]
        for thread in self.threads:
            thread.start()
            
    def getTask(self):
        #return self.taskQueue.get()
        return self.taskQueue.take()
    
    def putTask(self, runner, task):
        self.taskQueue.put((runner, task))
        
    def stop(self):
        self.putTask(self.kill, self.kill)
        for thread in self.threads:
            thread.join()
            
class PoolThread(Thread):
    def __init__(self, threadPool):
        Thread.__init__(self)
        self.pool = threadPool
    
    def run(self):
        while True:
            runner, task = self.pool.getTask()
            if runner == self.pool.kill:
                self.pool.putTask(self.pool.kill, self.pool.kill)
                break
            runner.process(task)